import os
import sys

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2,4"


project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
print(project_dir_path)
sys.path.append(project_dir_path)

from omics_bert_label_encoding import MakeLabelsUK
from omics_bert_label_encoding.label.check_label_files_uk import CheckLabelFilesUK
from omics_bert_label_encoding.label.processing_labels_uk import process_label_embeddings_uk


def main(dataset_path, save_path, id_file_path, encoding_file_path, instance_file_path, categories_file_path, api_key, expected_lines, regenerate, model_name, tensor_parallel_size):
    print("--- Step 1: Build Labels ---")
    build_labels = MakeLabelsUK(dataset_path, save_path, id_file_path, encoding_file_path,
                                         instance_file_path, categories_file_path, api_key)

    build_labels.main()

    print("\n--- Step 2: Check and Regenerate Labels ---")
    check_path = os.path.join(save_path, "labels")
    expected_lines = 336
    regenerate = True

    checker = CheckLabelFilesUK(
        check_path=check_path,
        expected_line_count=expected_lines,
        regenerate=regenerate,
        dataset_path=dataset_path,
        id_file_path=id_file_path,
        encoding_file_path=encoding_file_path,
        instance_file_path=instance_file_path,
        categories_file_path=categories_file_path,
        api_key=api_key
    )
    checker.main()

    print("\n--- Step 3: Process Label Embeddings ---")
    labels_dir = os.path.join(save_path, "labels")
    process_label_embeddings_uk(
        labels_json_path=labels_dir,
        output_dir=save_path,
        model_name=model_name,
        tensor_parallel_size=tensor_parallel_size
    )


if __name__ == '__main__':
    project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
    print(project_dir_path)
    sys.path.append(project_dir_path)

    dataset_path = "/home/xw/python3test/omics-bert-datasets/datasets/ready/ukbiobank/v4.2"
    save_path = os.path.join(project_dir_path, "embeddings_local/ukbiobank/v4-labels_embeddings")

    id_file_path = os.path.join(project_dir_path, "datasets/raw/ukbiobank/v3/id.csv")
    encoding_file_path = os.path.join(project_dir_path, "datasets/raw/ukbiobank/v3/encoding.csv")
    instance_file_path = os.path.join(project_dir_path, "datasets/raw/ukbiobank/v3/instance.csv")
    categories_file_path = os.path.join(project_dir_path, "datasets/raw/ukbiobank/v3/category.csv")

    api_key = ""
    expected_lines = 336
    regenerate = True
    model_name = "Qwen/Qwen3-Embedding-8B"
    tensor_parallel_size = 2

    main(dataset_path, save_path, id_file_path, encoding_file_path, instance_file_path, categories_file_path, api_key, expected_lines, regenerate, model_name, tensor_parallel_size)
